# Module aem.two_d.single.base -- base class for single-layer two-d elements
#
# Copyright (c) 2000, V.Kelson and M.Bakker
#           (c) 2001, Willem Jan Zaadnoordijk

import copy
import types
from math import sqrt
from numpy import *

import geometry.xy
XY = geometry.xy.XY

import aem.container
RowList = aem.container.RowList
ColumnList = aem.container.ColumnList
TraceCheck = aem.container.TraceCheck
import aem.matrix


class Element:
	'''Base class for analytic elements

	Attributes:

	- parent: parent of the element (often the aquifer)
	- parameters: strength parameters are stored as a ColumnList
	- solver: solver specified for the element (often None)
	- elementList: list of elements. For example the list of elements in the aquifer, or a list of linesinks in a linesink stringIs self by default'''
	def __init__(self,parent=None):
		'''Constructor'''
		if parent:
			self.parent = parent
			self.parent.addElement(self)
		else:
			self.parent = None
		self.parameters = ColumnList([])
		self.solver = None
		self.elementList = [ self ]
		self.givenElement = 1

	def addElement(self,el):
		'''Add element to elementList'''
		raise 'Must overload Element.addElement'

	def numberOfParameters(self):
		'''Returns number of strength parameters'''
		return len(self.parameters)

	# Influence functions -- Overload these for new elements
	def potentialInfluence(self,xy,t=0):
		'''Returns potential influence as RowList

		Returns the potential influence function(s) for the element at
		the point xy at time t, as an array. The return values are
		stored in a row vector, one entry per parameter.'''
		raise "Must overload Element.potentialInfluence()"
	def dischargeInfluence(self,xy,t=0):
		'''Returns the discharge influence as RowList

		Returns the discharge influence function(s) for the element at
		the point xy at time t, as a matrix, one column per unknown value.
		Each column contains two values, Qx and Qy, associated with the
		corresponding unknown in self.parameters.'''
		raise "Must overload Element.dischargeInfluence()"
	def normalDischargeInfluence(self,xy,orientation=0,t=0):
		'''Returns normal discharge influence as RowList

		Returns the component of discharge normal to the vector orientation
		at the point xy at time t.  The return values are stored in a row vector, one
		entry per parameter.'''
		raise "Must overload Element.normalDischargeInfluence()"
	def rechargeInfluence(self,xy,t=0):
		'''Returns recharge influence as RowList

		Returns the recharge influence function(s) for the element at
		the point xy at time t, as an array. The return values are
		stored in a row vector, one entry per unknown.'''
		return RowList([0.0])
	def fluxInfluence(self,xyBuf,t=0):
		'''Returns flux influence as RowList

		Returns the integrated flux influence function(s) for the element
		across the path mapped out in xyBuf, where xyBuf is a list of
		points defining the path.  The return values are
		stored in a row vector, one entry per unknown.'''
		raise "Must overload Element.fluxInfluence()"
	def extractionInfluence(self,t=0):
		'''Returns extraction influence as RowList

		Returns the total extraction rate influence function for the
		element(s) at time t.'''
		raise "Must overload Element.extractionInfluence()"

	# Contribution functions -- do not need to be overloaded
	def potentialContribution(self,xy,t=0):
		'''Returns potential at xy at time t

		Returns the potential due to 'self' at the point xy at time t.
		Result is returned as a list of length 1. [Phi]'''
		assert isinstance(xy,XY),"XY argument is required"
		return self.potentialInfluence(xy,t) * self.parameters
	def dischargeContribution(self,xy,t=0):
		'''Returns discharge at xy at time t

		Returns the discharge due to 'self' at the point xy at time t.
		Result is returned as a ColumnList'''
		assert isinstance(xy,XY),"XY argument is required"
		return self.dischargeInfluence(xy,t) * self.parameters
	def normalDischargeContribution(self,xy,orientation,t=0):
		'''Returns normal discharge at xy at time t

		Returns the component of discharge normal to the vector orientation
		at the point xy at time t.  Result is returned as a list of length 1'''
		assert isinstance(xy,XY),"XY location is required"
		assert isinstance(orientation,XY),"XY orientation is required"
		return self.normalDischargeInfluence(xy,t,orientation) * self.parameters
	def rechargeContribution(self,xy,t=0):
		'''Returns the recharge at xy at time t

		Returns the recharge due to 'self' at the point xy at time t.
		Result is returned as a list of length 1. '''
		assert isinstance(xy,XY),"XY argument is required"
		return self.rechargeInfluence(xy,t) * self.parameters
	def fluxContribution(self,xyBuf,t=0):
		'''Returns flux across line xyBuf at time t

		Returns the integrated flux due to 'self' across the path xyBuf at time t.
		Result is returned as a list of length 1.'''
		assert isinstance(xy,ListType),"List of XY arguments is required"
		return self.fluxInfluence(xyBuf,t) * self.parameters
	def extractionContribution(self,t=0):
		'''Returns exatraction due to element
		Returns the total extraction rate influence function for the
		element(s). Result is returned as a list of length 1'''
		return self.extractionInfluence(t) * self.parameters

	# Accumulated functions
	def potential(self,xy,t=0):
		'''Returns potential at point xy at time t

		For nearly all elements but Aquifer elements this is
		achieved by calling the 'parent' element. Must be overloaded
		for Aquifer element'''
		assert isinstance(xy,XY),"XY argument is required"
		return self.parent.potential(xy,t)
	def discharge(self,xy,t=0):
		'''Returns the discharge at point xy at time t.

		For nearly all elements but Aquifer elements this is
		achieved by calling the 'parent' element. Must be overloaded
		for Aquifer element'''
		assert isinstance(xy,XY),"XY argument is required"
		return self.parent.discharge(xy,t)
	def normalDischarge(self,xy,orientation,t=0):
		'''Returns the discharge at xy at time t normal to orientation

		For nearly all elements but Aquifer elements this is
		achieved by calling the 'parent' element. Must be overloaded
		for Aquifer element'''
		assert isinstance(xy,XY),"XY location is required"
		assert isinstance(orientation,XY),"XY orientation is required"
		return self.parent.normalDischarge(xy,orientation,t)
	def recharge(self,xy,t=0):
		'''Returns the recharge at point xy at time t.

		For nearly all elements but Aquifer elements this is
		achieved by calling the 'parent' element. Must be overloaded
		for Aquifer element'''
		assert isinstance(xy,XY),"XY argument is required"
		return self.parent.recharge(xy,t)
	def flux(self,xyBuf,t=0):
		'''Returns the flux across the path xyBuf at time t

		For nearly all elements but Aquifer elements this is
		achieved by calling the 'parent' element. Must be overloaded
		for Aquifer element'''
		assert isinstance(xy,ListType),"List of XY arguments is required"
		return self.parent.flux(xyBuf,t)
	def head(self,xy,t=0):
		'''Returns the head at point xy at time t.

		For nearly all elements but Aquifer elements this is
		achieved by calling the 'parent' element. Must be overloaded
		for Aquifer element'''
		assert isinstance(xy,XY),"XY argument is required"
		return self.parent.head(xy,t)
	def extraction(self,t=0):
		'''Returns the extraction rate.

		For nearly all elements but Aquifer elements this is
		achieved by calling the 'parent' element. Must be overloaded
		for Aquifer element'''
		return self.parent.extraction()
	def headToPotential(self,xy,head,t=0):
		'''Converts head to potential at xy at time t'''
		return self.parent.headToPotential(xy,head)
	def potentialToHead(self,xy,pot,t=0):
		'''Converts potential to head at xy at time t'''
		return self.parent.potentialToHead(xy,pot)
	def transmissivity(self,xy,t=0):
		'''Returns transmissivity at xy at time t'''
		return self.parent.transmissivity(xy,t)

	# Solution methods

	def solve(self,factor=1.0,elementList=None,t=0):
		'''Compute solution directly


		A solution is compiled to determine the change of the parameters
		with respect to the current value of the parameters. The factor
		may be set to less than one for relaxation of the solution
		Procedure:
		Compile full matrix for all elements in elementList
		Call aem.matrix.Solver(rows) for full solution
		Pass solution back to elements'''

		if not(self.givenElement):
			if elementList: solutionList = elementList
			else: solutionList = self.elementList
			rows = []
			columns = []
			for e in solutionList:
				mr,mc = e.getMatrixContribution(solutionList)
				rows = rows + mr
				columns = columns + mc
			s = aem.matrix.Solver(rows)
			x = s.solve()
			# Add solution to current value of parameter
			for i in range(len(x)):
				columns[i].element.parameters[columns[i].index] = \
					factor * x[i][0] + \
					columns[i].element.parameters[columns[i].index]

	def solveIteratively(self,factor=1.0,recursive=0,t=0):
		'''Compute solution iteratively'''
		for e in self.elementList:
			if recursive>1 and len(e.elementList)>1:
				e.solveIteratively(factor,recursive-1,t=t)
			else:
				e.solve(factor,e.elementList,t=t)

	def getMatrixCoefficients(self,fa,ctlPoints,t=0):
		'''Return matrix coefficients at ctlPoints for function fa

		Returns a list that contains matrix coefficients for 'self' at the
		point (or interval) in 'ctlPoints' '''
		return []
	def getMatrixContribution(self,others,t=0):
		'''Return matrix contribution

		Returns a tuple (mcr,mc) where mc is a list of MatrixColumn objects
		and mr is a list of MatrixRow objects'''
		return ([],[])
	def setParameter(self,index,value):
		'''Updates a strength parameter for the element

		Not used at this time'''
		assert index == 0,'Index out of range'
		self.parameters[index] = value
	def check(self,t=0):
		'''Checks the boundary condition

		Returns the difference between computed and specified boundary condition'''
		return 0.0

	# Methods for tracing
	def nearElement(self,xyz1,xyz2,step,t=0):
		'''Peforms something special when tracing near/over/into element

		Input: xyz1: last point of traceline, xyz2: candidate next point of trace line
		Returns: TraceCheck object'''
		rv = TraceCheck()
		return rv

	# Base class methods for computational testing
	def numericalGradient(self,xy,delta=0.001,t=0):
		'''
		Computes the gradient in the vicinity of 'xy', with dx=dy='delta' '''
		dx = XY(delta,0)
		dy = XY(0,delta)
		return XY( (self.potentialContribution(xy+dx,t) -
			    self.potentialContribution(xy-dx,t))/(2.0*delta),
			   (self.potentialContribution(xy+dy,t) -
			    self.potentialContribution(xy-dy,t))/(2.0*delta) )
	def numericalLaplacian(self,xy,delta=0.001,t=0):
		'''
		Computes the Laplacian in the vicinity of 'xy' with dx=dy='delta' '''
		dx = XY(delta,0)
		dy = XY(0,delta)
		return ( self.potentialContribution(xy+dx,t) + \
			 self.potentialContribution(xy-dx,t) + \
			 self.potentialContribution(xy+dy,t) + \
			 self.potentialContribution(xy-dy,t) - \
			 4.0*self.potentialContribution(xy,t) ) /(delta*delta)
